import os
from dotenv import load_dotenv, find_dotenv # 导入 find_dotenv 帮助定位
from langchain.chains import LLMChain
from langchain.chat_models import ChatOpenAI
from langchain.globals import set_verbose
from langchain.prompts import ChatPromptTemplate
from langchain.chains import SimpleSequentialChain
from langchain_core.runnables import RunnableSequence

# 加载 .env 文件中的环境变量 (增强调试)
load_dotenv(dotenv_path=find_dotenv(usecwd=True), verbose=True, override=True)

# 从环境变量加载 API 密钥和基础 URL
api_key = os.getenv("OPENAI_API_KEY")
api_base = os.getenv("OPENAI_API_BASE")
os.environ["OPENAI_API_KEY"] = api_key
os.environ["OPENAI_API_BASE"] = api_base

#顺序链 simpleSequentialChain 只支持固定的链路
chat_model = ChatOpenAI(
    temperature=0,
    model="gpt-3.5-turbo",
)

#chain 1
first_prompt = ChatPromptTemplate.from_template("帮我给{product}的公司起一个响亮容易记忆的名字?")
chain_one = LLMChain(
    llm=chat_model,
    prompt=first_prompt,
    verbose=True,
)

#chain 2
second_prompt = ChatPromptTemplate.from_template("用5个词来描述一下这个公司名字：{company_name}")
chain_two = LLMChain(
    llm=chat_model,
    prompt=second_prompt,
    verbose=True, #打开日志
    output_key="word"
)

#chain 3
third_prompt = ChatPromptTemplate.from_template("把{word}翻译成英文")
chain_three = LLMChain(
    llm=chat_model,
    prompt=third_prompt,
    verbose=True, #打开日志
)

# 把chain_one的结果传递给chain_two使用，得出最终结果
overall_simple_chain = SimpleSequentialChain(
    chains=[chain_one, chain_two, chain_three],
    verbose=True, #打开日志
)

overall_simple_chain.invoke("AI教育培训机构")